import os
import pickle
from collections import deque

import cv2
import numpy as np
import imutils


class ARProcessor:
    def __init__(self, marker_folder, calibrate=False, calibration_file='hj.mtl'):
        self.calibrate = calibrate
        self.calibration_file = calibration_file
        self.MTX = None
        self.DIST = None
        self.OBJPoints = deque(maxlen=5)
        self.IMGPoints = deque(maxlen=5)

        # 加载所有标记图像并提取SIFT特征
        self.markers = {}
        self.load_markers(marker_folder)

        # 初始化SIFT检测器
        self.sift = cv2.SIFT_create()
        if self.sift is None:
            raise RuntimeError(
                "无法创建 SIFT 检测器，请确认 OpenCV 是否包含 contrib 模块。\n"
                "尝试运行: pip install opencv-python-contrib"
            )

        # 设置FLANN匹配器
        FLANN_INDEX_KDTREE = 1
        index_params = dict(algorithm=FLANN_INDEX_KDTREE, trees=5)
        search_params = dict(checks=50)
        self.flann = cv2.FlannBasedMatcher(index_params, search_params)

        # 相机内参和畸变系数（如果未提供则使用默认值，实际应用中应通过相机标定获得）
        if not self.calibrate or not self._load_calibration():
            # 使用提供的默认相机内参和畸变系数
            self.MTX = np.array([
                [1.45570225e+03, 0.00000000e+00, 9.37301305e+02],
                [0.00000000e+00, 1.46024376e+03, 7.21115122e+02],
                [0.00000000e+00, 0.00000000e+00, 1.00000000e+00]
            ], dtype=np.float32)

            self.DIST = np.array([
                0.12919168, -0.34917103, -0.0009397, -0.00367339, 0.10893386
            ], dtype=np.float32)
        else:
            print("加载标定文件成功，使用文件中的相机参数")

        # 标记物的3D坐标（假设标记物在平面z=0上，单位为毫米）
        self.marker_size = 100  # 标记物边长
        self.marker_3d_points = np.array([
            [0, 0, 0],
            [self.marker_size, 0, 0],
            [self.marker_size, self.marker_size, 0],
            [0, self.marker_size, 0]
        ], dtype=np.float32)

    def load_markers(self, folder_path):
        """加载指定文件夹下的所有标记图像"""
        for filename in os.listdir(folder_path):
            if filename.endswith('.png') or filename.endswith('.jpg'):
                marker_name = os.path.splitext(filename)[0]
                marker_img = cv2.imread(os.path.join(folder_path, filename), 0)
                if marker_img is None:
                    raise FileNotFoundError(f"无法找到标记图像: {filename}")

                # 提取SIFT特征
                sift = cv2.SIFT_create()
                kp, des = sift.detectAndCompute(marker_img, None)
                self.markers[marker_name] = {
                    'image': marker_img,
                    'kp': kp,
                    'des': des
                }

    def process_frame(self, frame):
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        frame_kp, frame_des = self.sift.detectAndCompute(gray, None)

        if frame_des is None:
            return None, frame

        best_match_count = 0
        best_marker_name = None
        best_homography = None
        best_refined_corners = None

        MIN_MATCH_COUNT = 10  # 增加最低匹配点数

        for marker_name, marker_data in self.markers.items():
            marker_kp = marker_data['kp']
            marker_des = marker_data['des']

            # 特征匹配
            matches = self.flann.knnMatch(marker_des, frame_des, k=2)
            good_matches = [m for m, n in matches if m.distance < 0.7 * n.distance]

            print(f"{marker_name} 成功获取到角点数量: {len(good_matches)}")
            if len(good_matches) < MIN_MATCH_COUNT:
                continue  # 跳过不满足条件的标记

            # 计算单应性矩阵
            src_pts = np.float32([marker_kp[m.queryIdx].pt for m in good_matches]).reshape(-1, 1, 2)
            dst_pts = np.float32([frame_kp[m.trainIdx].pt for m in good_matches]).reshape(-1, 1, 2)
            homography, mask = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC, 5.0)

            if homography is None or np.isnan(homography).any():
                continue  # 跳过无效的单应性矩阵

            # 映射标记图像的四个角点到当前帧
            h_marker, w_marker = marker_data['image'].shape
            corners = np.float32([[0, 0], [0, h_marker - 1], [w_marker - 1, h_marker - 1], [w_marker - 1, 0]]).reshape(
                -1, 1, 2)
            transformed_corners = cv2.perspectiveTransform(corners, homography)

            # 检查角点是否在图像范围内
            h_frame, w_frame = gray.shape
            valid_corners = []
            for pt in transformed_corners.reshape(-1, 2):
                x, y = pt
                if 0 <= x < w_frame and 0 <= y < h_frame:
                    valid_corners.append([x, y])
                else:
                    print(f"检测到越界角点: ({x}, {y})")

            if len(valid_corners) != 4:
                print(f"{marker_name} 有效角点数量为 {len(valid_corners)}，不是 4 个，跳过该帧")
                continue  # 跳过无效的角点

            valid_corners = np.array(valid_corners, dtype=np.float32).reshape(-1, 1, 2)

            # 精确定位角点
            criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)
            refined_corners = cv2.cornerSubPix(gray, valid_corners, (11, 11), (-1, -1), criteria)

            # 验证角点间的相对距离是否合理
            def check_corner_distances(corners, expected_size):
                # 计算相邻角点之间的距离
                distances = [np.linalg.norm(corners[i] - corners[(i + 1) % 4]) for i in range(4)]
                avg_distance = sum(distances) / len(distances)
                # 检查每个边长是否接近预期大小
                tolerance = 0.6  # 允许的最大误差比例
                for d in distances:
                    if abs(d - expected_size) > expected_size * tolerance:
                        return False
                return True

            expected_size = self.marker_size  # 使用你的标记物的实际大小
            if not check_corner_distances(refined_corners, expected_size):
                print(f"{marker_name} 角点间距不符合预期，可能不是正确的标记物")
                continue  # 如果角点间距不符合预期，则跳过

            # 记录最佳匹配
            if len(good_matches) > best_match_count:
                best_match_count = len(good_matches)
                best_marker_name = marker_name
                best_homography = homography
                best_refined_corners = refined_corners

        if best_marker_name is None:
            print("未能找到任何有效的标记")
            return None, frame

        # 如果检测到足够多的匹配点，则进行相机标定, 否则返回None
        if self.calibrate:
            objp = np.zeros((2 * 2, 3), np.float32)
            objp[:, :2] = np.mgrid[0:2, 0:2].T.reshape(-1, 2)

            # 计算角点的空间坐标
            # 只有当角点位置变化较大时才记录
            if len(self.IMGPoints) == 0 or np.sum(np.abs(self.IMGPoints[-1] - best_refined_corners)) > 1e-3:
                self.OBJPoints.append(objp)
                self.IMGPoints.append(best_refined_corners)
            # 进行相机标定，返回相机内参矩阵和畸变系数
            if len(self.IMGPoints) >= 10:
                ret, mtx, dist, rvecs, tvecs = cv2.calibrateCamera(
                    list(self.OBJPoints), list(self.IMGPoints), gray.shape[::-1], None, None
                )
                if ret:
                    self.MTX = mtx
                    self.DIST = dist
                    self._save_calibration()
                    self.calibrate = False

                    _, rvec, tvec, _ = cv2.solvePnPRansac(self.marker_3d_points, best_refined_corners, self.MTX,
                                                          self.DIST)
                else:
                    print("标定失败，继续采集数据...")
                    return None, frame
            else:
                print(f"正在标定... 已采集 {len(self.IMGPoints)} 帧")
                return None, frame
        else:
            _, rvec, tvec, _ = cv2.solvePnPRansac(self.marker_3d_points, best_refined_corners, self.MTX, self.DIST)
            # print("相机内参矩阵: \n" + np.array2string(self.MTX, precision=4))
            # print("畸变系数: \n" + np.array2string(self.DIST, precision=4))
            print("标定测试结束")

        # 构建变换矩阵
        rotation_matrix, _ = cv2.Rodrigues(rvec)
        transformation_matrix = np.eye(4)
        transformation_matrix[:3, :3] = rotation_matrix
        transformation_matrix[:3, 3] = tvec.flatten()
        print("变换矩阵: \n" + np.array2string(transformation_matrix, precision=4))

        # 在返回前绘制矩形框
        frame_with_box = frame.copy()
        cv2.polylines(frame_with_box, [np.int32(best_refined_corners)], True, (0, 255, 0), 3, cv2.LINE_AA)

        # 绘制坐标轴
        axis_length = self.marker_size * 0.5  # 坐标轴长度为标记大小的一半
        axis_points = np.float32([
            [0, 0, 0],              # origin
            [axis_length, 0, 0],    # x-axis
            [0, axis_length, 0],    # y-axis
            [0, 0, -axis_length]    # z-axis
        ])

        # 投影3D坐标轴到图像平面
        imgpts, jac = cv2.projectPoints(axis_points, rvec, tvec, self.MTX, self.DIST)

        # 检查是否包含 NaN 或 Inf
        if not np.isfinite(imgpts).all():
            print("投影点包含非法值，跳过绘制坐标轴")
            return {best_marker_name: transformation_matrix}, frame_with_box

        imgpts = np.int32(imgpts).reshape(-1, 2)

        # 绘制坐标轴线条
        origin = tuple(imgpts[0])
        cv2.line(frame_with_box, origin, tuple(imgpts[1]), (0, 0, 255), 5)  # X轴 - 红色
        cv2.line(frame_with_box, origin, tuple(imgpts[2]), (255, 0, 0), 5)  # Y轴 - 蓝色
        cv2.line(frame_with_box, origin, tuple(imgpts[3]), (0, 255, 255), 5)  # Z轴 - 黄色

        print(f"{best_marker_name}标记检测成功")

        # 返回标记名称和变换矩阵
        return {best_marker_name: transformation_matrix}, frame_with_box

    def _load_calibration(self):
        try:
            with open(self.calibration_file, 'rb') as f:
                data = pickle.load(f)
                self.MTX = data['mtx']
                self.DIST = data['dist']
                return True
        except (FileNotFoundError, EOFError):
            return False

    def _save_calibration(self):
        data = {'mtx': self.MTX, 'dist': self.DIST}
        with open(self.calibration_file, 'wb') as f:
            pickle.dump(data, f)
